import numpy as np
from scipy.sparse.linalg import spsolve
from scipy.sparse import spdiags

class Inter_Solver:
    def __init__(self,A,b,tol=1e-8,x=0,maxiter=50):
        self.A = A
        self.b = b
        self.x = x
        self.tol = tol
        self.maxiter = maxiter

    def __str__(self):
        return "Faster slover!"

    def MINRES(self):
        'MINimum RESidual Algorithm'

        A = self.A
        x = self.x
        b = self.b
        tol = self.tol
        maxiter = self.maxiter

        if x != 0:
            b -= A@x

        
        k = 1
        beta_0 = np.sqrt(b@b); q = b/beta_0
        r_temp = A@q
        alpha = q@r_temp; v = r_temp - alpha*q; beta = np.sqrt(v@v)

        #print(q,alpha,beta)

        if beta == 0:
            x+=b/alpha
            return x
        else:
            q_new = v/beta
        # c_old 表 c_{k-1},  c_new 表 c_{k+1}, c 表 c_k
        c_old = 1; s_old = 0; p_old = 0
        c,s,gamma = self.Givens(alpha,beta)

        p = q/gamma; rho = -beta_0*s; tau = beta_0*c
        x+=tau*p
        

        while (np.abs(rho) > beta_0*tol)&(k<maxiter):
            r_temp = A@q_new
            alpha_new = q_new@r_temp; v = r_temp-alpha_new*q_new - beta*q; beta_new = np.sqrt(v@v)

            eplison_old = s_old*beta; hat_beta = c_old*beta
            delta = c*hat_beta+s*alpha_new; hat_alpha_new = -s*hat_beta+c*alpha_new

            c_new, s_new, gamma_new = self.Givens(hat_alpha_new,beta_new)

            tau_new = rho*c_new; rho_new = -rho*s_new
            p_new = (q_new-eplison_old*p_old-delta*p)/gamma_new

            x+=tau_new*p_new

            #update k = k + 1
            k+=1
            p_old = np.copy(p)
            s_old = np.copy(s)
            c_old = np.copy(c)

            q = np.copy(q_new)
            p = np.copy(p_new)
            beta = np.copy(beta_new)
            s = np.copy(s_new)
            c = np.copy(c_new)
            rho = np.copy(rho_new)

            if np.abs(beta) > 1e-15:
                q_new = v/beta
                #print(q_new@q_new)

            #print(np.abs(rho))
        print(k)
        return x




    def MINRES_Pre(self,M):
        'Preconditioner MINimum RESidual Algorithm'

        A = self.A
        x = self.x
        b = self.b
        tol = self.tol
        maxiter = self.maxiter

        if x != 0:
            b-= A@x

        
        k = 1
        #bar_bar_b = np.linalg.solve(M,b)
        bar_bar_b = spsolve(M,b)
        #print(np.max(np.abs(M@bar_bar_b-b)))
        beta_0 = np.sqrt(b@bar_bar_b); bar_bar_q = bar_bar_b/beta_0
        r_temp = A@bar_bar_q
        alpha = bar_bar_q@r_temp 
        #bar_bar_v = np.linalg.solve(M,r_temp) - alpha*bar_bar_q
        bar_bar_v = spsolve(M,r_temp) - alpha*bar_bar_q
        #print(np.max(np.abs(M@bar_bar_v-r_temp)))
        beta = np.sqrt(bar_bar_v@M@bar_bar_v)

        if beta == 0:
            x+=bar_bar_b/alpha
        else:
            bar_bar_q_new = bar_bar_v/beta

        c_old = 1; s_old = 0; bar_bar_p_old = 0
        c,s,gamma = self.Givens(alpha,beta)

        bar_bar_p = bar_bar_q/gamma; rho = -beta_0*s; tau = beta_0*c
        x+=tau*bar_bar_p


        while (np.abs(rho) > beta_0*tol)&(k<maxiter):
            r_temp = A@bar_bar_q_new
            alpha_new = bar_bar_q_new@r_temp
            #bar_bar_v = np.linalg.solve(M,r_temp) - alpha_new*bar_bar_q_new - beta*bar_bar_q;
            bar_bar_v = spsolve(M,r_temp) - alpha_new*bar_bar_q_new - beta*bar_bar_q 
            beta_new = np.sqrt(bar_bar_v@M@bar_bar_v)


            eplison_old = s_old*beta; hat_beta = c_old*beta
            delta = c*hat_beta+s*alpha_new; hat_alpha_new = -s*hat_beta+c*alpha_new

            c_new, s_new, gamma_new = self.Givens(hat_alpha_new,beta_new)

            tau_new = rho*c_new; rho_new = -rho*s_new
            bar_bar_p_new = (bar_bar_q_new - eplison_old*bar_bar_p_old-delta*bar_bar_p)/gamma_new

            x+=tau_new*bar_bar_p_new

            #update k = k+1
            k+=1
            bar_bar_p_old = np.copy(bar_bar_p)
            s_old = np.copy(s)
            c_old = np.copy(c)

            bar_bar_q = np.copy(bar_bar_q_new)
            bar_bar_p = np.copy(bar_bar_p_new)
            beta = np.copy(beta_new)
            s = np.copy(s_new)
            c = np.copy(c_new)
            rho = np.copy(rho_new)

            if np.abs(beta) > 1e-15:
                bar_bar_q_new = bar_bar_v/beta

            #print(np.abs(rho))


        print(k)
        return x



    


    def Givens(self,a,b):
        if np.abs(b) < 1e-15:
            c = 1.0; s = 0.0

        if np.abs(b)>np.abs(a):
            tau = a/b; s = 1.0/np.sqrt(1+tau**2); c = s*tau 
        
        else:
            tau = b/a; c = 1.0/np.sqrt(1+tau**2); s = c*tau

        r = c*a+s*b

        #print(s**2+c**2-1.0,-s*a+c*b)

        return c,s,r
        





        

        






if __name__ == '__main__':
    import numpy as np
    from scipy.sparse import csr_matrix
    from scipy.sparse.linalg import minres, gmres


    n = 1000000
    maxiter = 50
    if False:
        A = np.random.random((n,n))
        b = np.random.random(n)
        A = A + A.T
        #M = np.eye(n)
        Lam, Q = np.linalg.eig(A)
        #M = np.diag(np.abs(np.diag(A))+1.0)
        M = np.diag(np.abs(Lam))
    else:
        ii = np.zeros(3*n-2,dtype=int)
        jj = np.arange(3*n-2,dtype=int)
        d = -np.ones(3*n-2,dtype=float)

        ii[:n] = np.arange(n,dtype=int)
        jj[:n] = np.arange(n,dtype=int)
        d[:n] = np.arange(n,dtype=float)+1

        ii[n:2*n-1] = np.arange(n-1,dtype=int)+1
        jj[n:2*n-1] = np.arange(n-1,dtype=int)


        ii[2*n-1:] = np.arange(n-1,dtype=int)
        jj[2*n-1:] = np.arange(n-1,dtype=int)+1


        print('ok1')
        A = csr_matrix((d.flat,(ii.flat,jj.flat)),shape=(n,n))
        print('ok2')
        b = np.zeros(n,dtype=float)+1.0
        print('ok3')
        #b = A@b
        print('ok4')
        ii = np.arange(n,dtype=int)
        dd = np.arange(n,dtype=float)+1.0
        M = csr_matrix((dd.flat,(ii.flat,ii.flat)),shape=(n,n))
        print('ok5')
    print('ok6')
    #C_inv = np.linalg.inv(np.sqrt(M))
    #Lam, Q = np.linalg.eig(A)
    #Lam1, Q1 = np.linalg.eig(C_inv@A@C_inv)
    #print(np.max(np.abs(Lam))/np.min(np.abs(Lam)),np.max(np.abs(Lam1))/np.min(np.abs(Lam1)))


    #x = np.linalg.solve(A,b)
    #r = A@x-b
    #print(np.sqrt(r@r))
    #x, exitCode = minres(A,b,maxiter=maxiter,tol=1e-15)
    #x, exitCode = gmres(A,b,M=M,maxiter=maxiter,tol=1e-15)
    #r = A@x-b
    #print(np.sqrt(r@r),exitCode)


    ISr = Inter_Solver(A,b,maxiter=maxiter,tol=1e-15)
    #xs = ISr.MINRES()
    #r = A@xs-b
    #print(np.sqrt(r@r))

    xss = ISr.MINRES_Pre(M)
    r = A@xss-b
    print(np.sqrt(r@r))



